# coding=utf-8
##
## Author: jmdvirus@aliyun.com
##
## Create: 2019年02月12日 星期二 18时07分43秒
##

import numpy as np
import os
import gzip
import struct
import logging
import mxnet as mx
import matplotlib.pyplot as plt

logging.getLogger().setLevel(logging.DEBUG)

def read_data(label_url, image_url):
    with gzip.open(label_url) as flbl:
        magic,num = struct.unpack(">II", flbl.read(8))
        label = np.fromstring(flbl.read(), dtype = np.int8)
    with gzip.open(image_url, 'rb') as fimg:
        magic,num,rows,cols = struct.unpack(">IIII", fimg.read(16))
        image = np.fromstring(fimg.read(), dtype=np.int8)
        image = image.reshape(len(label), 1, rows, cols)
        image = image.astype(np.float32)/255.0
    return (label, image)

data_path = "/opt/data/data/av.data.mnist/"
(train_lbl, train_img) = read_data(data_path + 'train-labels-idx1-ubyte.gz', data_path + 'train-images-idx3-ubyte.gz')

(val_lbl, val_img) = read_data(data_path + 't10k-labels-idx1-ubyte.gz', data_path + 't10k-images-idx3-ubyte.gz')

batch_size = 32

train_iter = mx.io.NDArrayIter(train_img, train_lbl, batch_size, shuffle = True)
val_iter = mx.io.NDArrayIter(val_img, val_lbl, batch_size)

#for i in range(10):
#    plt.subplot(1, 10, i+1)
#    plt.imshow(train_img[i].reshape(28,28), cmap='Greys_r')
#    plt.axis('off')
#plt.show()
print('label: %s' % (train_lbl[0:10],))

data = mx.symbol.Variable('data')

def conv_act(indata, layer, kernel, num):
    conv = mx.sym.Convolution(data = indata, name="conv" + layer, kernel=kernel, num_filter=num)
    bn = mx.sym.BatchNorm(data = conv, name="bn" + layer, fix_gamma=False)
    act = mx.sym.Activation(data=bn, name="act"+layer, act_type="relu")
    return act

net = conv_act(data, "1", (5,5), 32)
pool1 = mx.sym.Pooling(data=net, name='pool1', pool_type='max', 
        kernel=(3,3), stride=(2,2))

net = conv_act(pool1, "2", (5,5), 64)
pool2 = mx.sym.Pooling(data=net, name='pool2', pool_type='max', kernel=(3,3),
        stride=(2,2))

conv3 = mx.sym.Convolution(data=pool2, name='conv3', kernel=(3,3), num_filter=10)
pool3 = mx.sym.Pooling(data=conv3, name='pool3', global_pool=True, pool_type='avg',
    kernel=(1,1))

flatten = mx.sym.Flatten(data=pool3, name='flatten')
net = mx.sym.SoftmaxOutput(data=flatten, name='softmax')

shape = { "data" : (batch_size, 1, 28, 28 ) }

def plot_network():
    mx.viz.print_summary(symbol=net, shape=shape)
    mx.viz.plot_network(symbol=net, shape=shape).view()

def module_train(data):
    module = mx.mod.Module(symbol=data, context=mx.cpu(0))
    module.fit(
            train_iter,
            eval_data = val_iter,
            optimizer = 'sgd',
            optimizer_params = {'learning_rate':0.2,
                'lr_scheduler' : mx.lr_scheduler.FactorScheduler(step=60000/batch_size, factor=0.9) },
            num_epoch = 20,
            batch_end_callback = mx.callback.Speedometer(batch_size, 60000/batch_size)
            )

if __name__ == "__main__":
    plot_network()
    # module_train(net)

